#include "avltree.h"
#include "mem.h"

typedef struct z_tree_node z_tree_node;

/* @internal 
 * @breif Tree
 */
struct z_tree
{
  size_t nnodes;                /* number of nodes */
  z_tree_node *root;            /* root node */
  z_compare_func *key_compare_func; /* key compare */
  z_destroy_notify *key_destroy_func; /* key destory */
  z_destroy_notify *value_destroy_func; /* value destroy */
};

struct z_tree_node
{
  void *key;                    /* key */
  void *value;                  /* value */
  z_tree_node* left;            /* the left child */
  z_tree_node* right;           /* the right child */
  int balance, height;
};

z_tree* z_tree_new(z_compare_func key_compare_func)
{
  return z_tree_new_full(key_compare_func, NULL, NULL);
}

z_tree* z_tree_new_full(z_compare_func key_compare_func, z_destroy_notify key_destroy_func,
                        z_destroy_notify value_destroy_func)
{
  z_tree *tree;
  
  tree = z_new(z_tree, 1);
  tree->nnodes = 0;
  tree->key_compare_func = key_compare_func;
  tree->key_destroy_func = key_destroy_func;
  tree->value_destroy_func = value_destroy_func;
  
  tree->root = NULL;
  return tree;
}

static void z_tree_node_destroy(z_tree *tree, z_tree_node *node)
{
  if (!node)
    return;
  z_tree_node_destroy(tree, node->left);
  z_tree_node_destroy(tree, node->right);
  
  if (tree->key_destroy_func)
    tree->key_destroy_func(node->key);
  if (tree->value_destroy_func)
    tree->value_destroy_func(node->value);
  
  zfree(node);
}

void z_tree_destory(z_tree *tree)
{
  z_tree_node_destroy(tree, tree->root);
  zfree(tree);
}

static z_tree_node* z_tree_node_new(void *key, void *value)
{
  z_tree_node *tree_node;
  
  tree_node = z_new(z_tree_node, 1);
  
  tree_node->key = key;
  tree_node->value = value;
  tree_node->left = tree_node->right = NULL;
  tree_node->balance = 0;
  tree_node->height = 1;
  return tree_node;
}

/* @internal
 * @breif caculate node balance and height
 */
static void z_tree_node_calculate(z_tree_node *node)
{
  int l_height = 0;
  int r_height = 0;
  if (node->left)
    l_height = node->left->height;
  if (node->right)
    r_height = node->right->height;
  node->height = (l_height > r_height ? l_height : r_height) + 1;
  node->balance = r_height - l_height;
}

/* Rotate left 
 */
static z_tree_node* z_tree_left_rotate(z_tree_node *node)
{
  z_tree_node *u = node;
  z_tree_node *v = u->right;
  z_tree_node *b = v->left;
  
  u->right = b;
  v->left = u;
  z_tree_node_calculate(u);
  z_tree_node_calculate(v);
  return v;
}

/* Right Rotate
 */
static z_tree_node* z_tree_right_rotate(z_tree_node *node)
{
  z_tree_node *u = node;
  z_tree_node *v = u->left;
  z_tree_node *b = v->right;
  
  u->left = b;
  v->right = u;
  z_tree_node_calculate(u);
  z_tree_node_calculate(v);
  return v;
}

/* double Left rotate in fact, this is combined by two
 * rotate:left rotate on lower part, the right rotate on upper
 * part,and this shuold reduce node calculation by once!
 */

static z_tree_node* z_tree_double_left(z_tree_node *node)
{
  z_tree_node *u = node;
  z_tree_node *v = v->right;
  z_tree_node *z = v->left;
  z_tree_node *b = z->left;
  z_tree_node *c = z->right;
  
  z->left = u;
  z->right = v;
  u->right = b;
  v->left = c;
  z_tree_node_calculate(u);
  z_tree_node_calculate(v);
  z_tree_node_calculate(z);
  
  return z;
}

/* double right rotate
 */
static z_tree_node* z_tree_double_right(z_tree_node *node)
{
  z_tree_node *u = node;
  z_tree_node *v = u->left;
  z_tree_node *z = v->right;
  z_tree_node *b = z->right;
  z_tree_node *c = z->left;
  
  z->right = u;
  z->left = v;
  u->left = b;
  v->right = c;
  
  z_tree_node_calculate(u);
  z_tree_node_calculate(v);
  z_tree_node_calculate(z);
  return z;
}

static z_tree_node* z_tree_node_rotate(z_tree_node *node)
{
  z_tree_node_calculate(node);
  
  switch(node->balance) {
  case -2:
    switch(node->left->balance) {
    case -1:
    case 0:
      return z_tree_right_rotate(node);
    case 1:
      return z_tree_double_right(node);
    }
  case -1:
  case 0:
  case 1:
    return node;
  case 2:
    switch(node->right->balance) {
    case -1:
      return z_tree_double_left(node);
    case 0:
    case 1:
      return z_tree_left_rotate(node);
    }
  }
  zerror1("Internnal error");
  return node;
}

static z_tree_node* z_tree_node_insert(z_tree *tree, z_tree_node *node, void *key, void *value)
{
  int cmp = tree->key_compare_func(key, node->key);
  
  if (cmp == 0) {
    node->value = value;
  } else if (cmp < 0) {
    if (node->left)
      node->left = z_tree_node_insert(tree, node->left, key, value);
    else {
      node->left = z_tree_node_new(key, value);
      tree->nnodes++;
    }
  } else {
    if (node->right)
      node->right = z_tree_node_insert(tree, node->right, key, value);
    else {
      node->right = z_tree_node_new(key, value);
      tree->nnodes++;
    }
  }
  return z_tree_node_rotate(node);
}

void z_tree_insert(z_tree *tree, void *key, void *value)
{
  if(!tree->root) {
    tree->root = z_tree_node_new(key, value);
    tree->nnodes++;
  } else {
    tree->root = z_tree_node_insert(tree, tree->root, key, value);
  }
}

static void z_tree_node_foreach(z_tree_node *node, z_hfunc *func, void *userdata)
{
  if (!node)
    return ;
  z_tree_node_foreach(node->left, func, userdata);
  func(node->key, node->value, userdata);
  z_tree_node_foreach(node->right, func, userdata);
}

void z_tree_foreach(z_tree *tree, z_hfunc *func, void *userdata)
{
  z_tree_node_foreach(tree->root, func, userdata);
}

uint32_t z_tree_size(z_tree *tree)
{
  return tree->nnodes;
}

static void* z_tree_node_lookup(z_tree *tree, z_tree_node *node, const void *key) 
{
  int cmp;
  if (!node)
    return NULL;
  cmp = tree->key_compare_func(node->key, key);
  if (cmp == 0)
    return node->value;
  else if (cmp < 0)
    return z_tree_node_lookup(tree, node->left, key);
  else
    return z_tree_node_lookup(tree, node->right, key);
}

void *z_tree_lookup(z_tree *tree, const void *key)
{
  return z_tree_node_lookup(tree, tree->root, key);
}

#ifdef DEBUGAVL
#include <string.h>
#include <stdio.h>  

int cmp(const void *a,const void *b)
{
  return strcmp((char *)a, (char *)b);
}

void print (void *key, void *value, void *userdata) 
{
  printf("key: %s-->value:%s\n", (char *)key, (char *)value);
}

int main() 
{
  z_tree *t = z_tree_new(cmp);

  z_tree_insert(t, "a", "kk");
  z_tree_insert(t, "a", "nn");
  z_tree_insert(t, "d", "qq");
  z_tree_insert(t, "e", "vv");
  z_tree_insert(t, "a", "ww");
  
  z_tree_foreach(t, print, NULL);
  
  z_tree_destory(t);
  return 0;
}
#endif
